// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <type_traits>
#include <nebula/array/array_base.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/compute/api_scalar.h>
#include <nebula/compute/api_vector.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/kernels/base_arithmetic_internal.h>
#include <nebula/compute/kernels/codegen_internal.h>
#include <nebula/compute/kernels/common_internal.h>

#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <nebula/core/visit_type_inline.h>

namespace nebula::compute::internal {

    namespace {

        template<typename OptionsType>
        struct CumulativeOptionsWrapper : public OptionsWrapper<OptionsType> {
            using State = CumulativeOptionsWrapper<OptionsType>;

            explicit CumulativeOptionsWrapper(OptionsType options)
                    : OptionsWrapper<OptionsType>(std::move(options)) {}

            static turbo::Result<std::unique_ptr<KernelState>> init(KernelContext *ctx,
                                                                    const KernelInitArgs &args) {
                auto options = turbo::checked_cast<const OptionsType *>(args.options);
                if (!options) {
                    return turbo::invalid_argument_error(
                            "Attempted to initialize KernelState from null FunctionOptions");
                }

                const auto &start = options->start;

                // Ensure `start` option, if given, matches input type
                if (start.has_value() && !start.value()->type->equals(*args.inputs[0])) {
                    TURBO_MOVE_OR_RAISE(auto casted_start,
                                        Cast(Datum(start.value()), args.inputs[0],
                                             CastOptions::Safe(), ctx->exec_context()));
                    auto new_options = OptionsType(casted_start.scalar(), options->skip_nulls);
                    return std::make_unique<State>(new_options);
                }
                return std::make_unique<State>(*options);
            }
        };

// The cumulative value is computed based on a simple arithmetic binary op
// such as Add, Mul, Min, Max, etc.
        template<typename Op, typename ArgType>
        struct CumulativeBinaryOp {
            using OutType = ArgType;
            using OutValue = typename GetOutputType<OutType>::T;
            using ArgValue = typename GetViewType<ArgType>::T;

            OutValue current_value;

            CumulativeBinaryOp() { current_value = Identity<Op>::template value<OutValue>; }

            explicit CumulativeBinaryOp(const std::shared_ptr<Scalar> start) {
                current_value = UnboxScalar<OutType>::Unbox(*start);
            }

            OutValue Call(KernelContext *ctx, ArgValue arg, turbo::Status *st) {
                current_value =
                        Op::template Call<OutValue, ArgValue, ArgValue>(ctx, arg, current_value, st);
                return current_value;
            }
        };

        template<typename ArgType>
        struct CumulativeMean {
            using OutType = Fp64Type;
            using ArgValue = typename GetViewType<ArgType>::T;
            int64_t count = 0;
            double sum = 0;

            CumulativeMean() = default;

            // start value is ignored for CumulativeMean
            explicit CumulativeMean(const std::shared_ptr<Scalar> start) {}

            double Call(KernelContext *ctx, ArgValue arg, turbo::Status *st) {
                sum += static_cast<double>(arg);
                ++count;
                return sum / count;
            }
        };

// The driver kernel for all cumulative compute functions.
// ArgType and OutType are the input and output types, which will
// normally be the same (e.g. the cumulative sum of an array of Int64Type will result in
// an array of Int64Type) with the exception of CumulativeMean, which will always return
// a double.
        template<typename ArgType, typename CumulativeState>
        struct Accumulator {
            using OutType = typename CumulativeState::OutType;
            using ArgValue = typename GetViewType<ArgType>::T;

            KernelContext *ctx;
            CumulativeState current_state;
            bool skip_nulls;
            bool encountered_null = false;
            NumericBuilder<OutType> builder;

            explicit Accumulator(KernelContext *ctx) : ctx(ctx), builder(ctx->memory_pool()) {}

            turbo::Status Accumulate(const ArraySpan &input) {
                turbo::Status st = turbo::OkStatus();

                if (skip_nulls || (input.get_null_count() == 0 && !encountered_null)) {
                    VisitArrayValuesInline<ArgType>(
                            input,
                            [&](ArgValue v) { builder.unsafe_append(current_state.Call(ctx, v, &st)); },
                            [&]() { builder.unsafe_append_null(); });
                } else {
                    int64_t nulls_start_idx = 0;
                    VisitArrayValuesInline<ArgType>(
                            input,
                            [&](ArgValue v) {
                                if (!encountered_null) {
                                    builder.unsafe_append(current_state.Call(ctx, v, &st));
                                    ++nulls_start_idx;
                                }
                            },
                            [&]() { encountered_null = true; });

                    TURBO_RETURN_NOT_OK(builder.append_nulls(input.length - nulls_start_idx));
                }

                return st;
            }
        };

        template<typename ArgType, typename CumulativeState, typename OptionsType>
        struct CumulativeKernel {
            using OutType = typename CumulativeState::OutType;
            using OutValue = typename GetOutputType<OutType>::T;

            static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
                const auto &options = CumulativeOptionsWrapper<OptionsType>::Get(ctx);
                Accumulator<ArgType, CumulativeState> accumulator(ctx);
                if (options.start.has_value()) {
                    accumulator.current_state = CumulativeState(options.start.value());
                } else {
                    accumulator.current_state = CumulativeState();
                }
                accumulator.skip_nulls = options.skip_nulls;

                TURBO_RETURN_NOT_OK(accumulator.builder.Reserve(batch.length));
                TURBO_RETURN_NOT_OK(accumulator.Accumulate(batch[0].array));

                std::shared_ptr<ArrayData> result;
                TURBO_RETURN_NOT_OK(accumulator.builder.finish_internal(&result));
                out->value = std::move(result);
                return turbo::OkStatus();
            }
        };

        template<typename ArgType, typename CumulativeState, typename OptionsType>
        struct CumulativeKernelChunked {
            using OutType = typename CumulativeState::OutType;
            using OutValue = typename GetOutputType<OutType>::T;

            static turbo::Status Exec(KernelContext *ctx, const ExecBatch &batch, Datum *out) {
                const auto &options = CumulativeOptionsWrapper<OptionsType>::Get(ctx);
                Accumulator<ArgType, CumulativeState> accumulator(ctx);
                if (options.start.has_value()) {
                    accumulator.current_state = CumulativeState(options.start.value());
                } else {
                    accumulator.current_state = CumulativeState();
                }
                accumulator.skip_nulls = options.skip_nulls;

                const ChunkedArray &chunked_input = *batch[0].chunked_array();
                TURBO_RETURN_NOT_OK(accumulator.builder.Reserve(chunked_input.length()));
                std::vector<std::shared_ptr<Array>> out_chunks;
                for (const auto &chunk: chunked_input.chunks()) {
                    TURBO_RETURN_NOT_OK(accumulator.Accumulate(*chunk->data()));
                }
                std::shared_ptr<ArrayData> result;
                TURBO_RETURN_NOT_OK(accumulator.builder.finish_internal(&result));
                out->value = std::move(result);
                return turbo::OkStatus();
            }
        };

        const FunctionDoc cumulative_sum_doc{
                "Compute the cumulative sum over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative sum computed over `values`. Results will wrap around on\n"
                 "integer overflow. Use function \"cumulative_sum_checked\" if you want\n"
                 "overflow to return an error. The default start is 0."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_sum_checked_doc{
                "Compute the cumulative sum over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative sum computed over `values`. This function returns an error\n"
                 "on overflow. For a variant that doesn't fail on overflow, use\n"
                 "function \"cumulative_sum\". The default start is 0."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_prod_doc{
                "Compute the cumulative product over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative product computed over `values`. Results will wrap around on\n"
                 "integer overflow. Use function \"cumulative_prod_checked\" if you want\n"
                 "overflow to return an error. The default start is 1."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_prod_checked_doc{
                "Compute the cumulative product over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative product computed over `values`. This function returns an error\n"
                 "on overflow. For a variant that doesn't fail on overflow, use\n"
                 "function \"cumulative_prod\". The default start is 1."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_max_doc{
                "Compute the cumulative max over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative max computed over `values`. The default start is the minimum\n"
                 "value of input type (so that any other value will replace the\n"
                 "start as the new maximum)."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_min_doc{
                "Compute the cumulative min over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative min computed over `values`. The default start is the maximum\n"
                 "value of input type (so that any other value will replace the\n"
                 "start as the new minimum)."),
                {"values"},
                "CumulativeOptions"};

        const FunctionDoc cumulative_mean_doc{
                "Compute the cumulative mean over a numeric input",
                ("`values` must be numeric. Return an array/chunked array which is the\n"
                 "cumulative mean computed over `values`. CumulativeOptions::start_value is \n"
                 "ignored."),
                {"values"},
                "CumulativeOptions"};

// Kernel factory for complex stateful computations.
        template<template<typename ArgType> typename State, typename OptionsType>
        struct CumulativeStatefulKernelFactory {
            VectorKernel kernel;

            CumulativeStatefulKernelFactory() {
                kernel.can_execute_chunkwise = false;
                kernel.null_handling = NullHandling::type::COMPUTED_NO_PREALLOCATE;
                kernel.mem_allocation = MemAllocation::type::NO_PREALLOCATE;
                kernel.init = CumulativeOptionsWrapper<OptionsType>::init;
            }

            template<typename Type>
            enable_if_number<Type, turbo::Status> Visit(const Type &type) {
                kernel.signature = KernelSignature::create(
                        {type.get_shared_ptr()},
                        OutputType(TypeTraits<typename State<Type>::OutType>::type_singleton()));
                kernel.exec = CumulativeKernel<Type, State<Type>, OptionsType>::Exec;
                kernel.exec_chunked = CumulativeKernelChunked<Type, State<Type>, OptionsType>::Exec;
                return turbo::OkStatus();
            }

            turbo::Status Visit(const DataType &type) {
                return turbo::unimplemented_error("Cumulative kernel not implemented for type ",
                                                  type.to_string());
            }

            turbo::Result<VectorKernel> create(const DataType &type) {
                TURBO_RETURN_NOT_OK(visit_type_inline(type, this));
                return kernel;
            }
        };

        template<template<typename ArgType> typename State, typename OptionsType>
        void MakeVectorCumulativeStatefulFunction(FunctionRegistry *registry,
                                                  const std::string func_name,
                                                  const FunctionDoc doc) {
            static const OptionsType kDefaultOptions = OptionsType::defaults();
            auto func =
                    std::make_shared<VectorFunction>(func_name, Arity::Unary(), doc, &kDefaultOptions);

            std::vector<std::shared_ptr<DataType>> types;
            types.insert(types.end(), numeric_types().begin(), numeric_types().end());

            CumulativeStatefulKernelFactory<State, OptionsType> kernel_factory;
            for (const auto &ty: types) {
                auto kernel = kernel_factory.create(*ty).value_or_die();
                KCHECK_OK(func->add_kernel(std::move(kernel)));
            }

            KCHECK_OK(registry->add_function(std::move(func)));
        }

// A kernel factory that forwards to CumulativeBinaryOp<Op, ...> for the given type.
// Need to use a struct because template-using declarations cannot appear in
// function scope.
        template<typename Op, typename OptionsType>
        struct MakeVectorCumulativeBinaryOpFunction {
            template<typename ArgType>
            using State = CumulativeBinaryOp<Op, ArgType>;

            static void Call(FunctionRegistry *registry, std::string func_name, FunctionDoc doc) {
                MakeVectorCumulativeStatefulFunction<State, OptionsType>(
                        registry, std::move(func_name), std::move(doc));
            }
        };

    }  // namespace

    void register_vector_cumulative_sum(FunctionRegistry *registry) {
        MakeVectorCumulativeBinaryOpFunction<Add, CumulativeOptions>::Call(
                registry, "cumulative_sum", cumulative_sum_doc);
        MakeVectorCumulativeBinaryOpFunction<AddChecked, CumulativeOptions>::Call(
                registry, "cumulative_sum_checked", cumulative_sum_checked_doc);

        MakeVectorCumulativeBinaryOpFunction<Multiply, CumulativeOptions>::Call(
                registry, "cumulative_prod", cumulative_prod_doc);
        MakeVectorCumulativeBinaryOpFunction<MultiplyChecked, CumulativeOptions>::Call(
                registry, "cumulative_prod_checked", cumulative_prod_checked_doc);

        MakeVectorCumulativeBinaryOpFunction<Min, CumulativeOptions>::Call(
                registry, "cumulative_min", cumulative_min_doc);
        MakeVectorCumulativeBinaryOpFunction<Max, CumulativeOptions>::Call(
                registry, "cumulative_max", cumulative_max_doc);

        MakeVectorCumulativeStatefulFunction<CumulativeMean, CumulativeOptions>(
                registry, "cumulative_mean", cumulative_max_doc);
    }

}  // namespace nebula::compute::internal
